import os
from PyPDF2 import PdfReader, PdfWriter
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter

# 创建水印 PDF
def create_watermark(content, output_path):
    c = canvas.Canvas(output_path, pagesize=letter)
    width, height = letter
    c.setFont("Helvetica", 50)
    c.setFillColorRGB(0.5, 0.5, 0.5, 0.5)  # 半透明的灰色
    c.rotate(45)
    c.drawString(100, height / 2, content)
    c.save()

# 添加水印到 PDF
def add_watermark(input_path, output_path, watermark_path):
    pdf_reader = PdfReader(input_path)
    pdf_writer = PdfWriter()
    watermark_reader = PdfReader(watermark_path)
    watermark_page = watermark_reader.pages[0]
    
    for page in pdf_reader.pages:
        page.merge_page(watermark_page)
        pdf_writer.add_page(page)
    
    with open(output_path, 'wb') as output_file:
        pdf_writer.write(output_file)

# 主函数
def main():
    directory = r'D:\软件安装包及其它'  
    pdf_files = find_pdf_files(directory)
    
    for pdf_file in pdf_files:
        watermark_path = 'watermark.pdf'
        create_watermark('Confidential', watermark_path)
        
        watermarked_pdf_file = pdf_file.replace('.pdf', '_watermarked.pdf')
        add_watermark(pdf_file, watermarked_pdf_file, watermark_path)
        print(f"添加水印并保存了 {pdf_file} 到 {watermarked_pdf_file}")

if __name__ == "__main__":
    main()